package com.example.drcbackend.util;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.example.drcbackend.entity.dto.ForeignKeyInfo;
import com.example.drcbackend.entity.dto.TableSchema;

import java.util.ArrayList;
import java.util.List;

public class SqlParserUtil {

    public static List<TableSchema> extractAllSchemas(String sql) {
        SQLStatementParser parser = new MySqlStatementParser(sql);
        List<SQLStatement> stmts = parser.parseStatementList();
        List<TableSchema> schemas = new ArrayList<>();

        for (SQLStatement stmt : stmts) {
            if (stmt instanceof SQLCreateTableStatement) {
                SQLCreateTableStatement create = (SQLCreateTableStatement) stmt;
                TableSchema schema = parseCreateStatement(create);
                schemas.add(schema);
            }
        }

        printSchemas(schemas);  // 调试输出
        return schemas;
    }

    private static TableSchema parseCreateStatement(SQLCreateTableStatement create) {
        TableSchema schema = new TableSchema();
        schema.setName(create.getName().getSimpleName());

        List<String> primaryKeys = new ArrayList<>();
        List<ForeignKeyInfo> foreignKeys = new ArrayList<>();
        List<String> uniqueColumns = new ArrayList<>();
        List<String> attributes = new ArrayList<>(); // 新增属性收集

        for (SQLTableElement element : create.getTableElementList()) {
            if (element instanceof SQLPrimaryKey) {
                SQLPrimaryKey pk = (SQLPrimaryKey) element;
                for (SQLSelectOrderByItem col : pk.getColumns()) {
                    primaryKeys.add(extractColumnName(col.getExpr()));
                }
            } else if (element instanceof SQLForeignKeyConstraint) {
                SQLForeignKeyConstraint fk = (SQLForeignKeyConstraint) element;
                String refTable = fk.getReferencedTableName().getSimpleName();
                List<SQLName> srcCols = fk.getReferencingColumns();
                List<SQLName> refCols = fk.getReferencedColumns();

                for (int i = 0; i < srcCols.size(); i++) {
                    foreignKeys.add(new ForeignKeyInfo(
                            srcCols.get(i).getSimpleName(),
                            refTable,
                            refCols.get(i).getSimpleName()
                    ));
                }
            } else if (element instanceof SQLUnique) {
                SQLUnique unique = (SQLUnique) element;
                for (SQLSelectOrderByItem col : unique.getColumns()) {
                    uniqueColumns.add(extractColumnName(col.getExpr()));
                }
            } else if (element instanceof SQLColumnDefinition) {
                SQLColumnDefinition column = (SQLColumnDefinition) element;
                String columnName = column.getName().getSimpleName();
                attributes.add(columnName); // 收集列名

                if (isPrimaryKeyColumn(column)) {
                    primaryKeys.add(columnName);
                }
                if (isUniqueColumn(column)) {
                    uniqueColumns.add(columnName);
                }
            }
        }

        schema.setPrimaryKeyColumns(primaryKeys);
        schema.setForeignKeyColumns(foreignKeys);
        schema.setUniqueColumns(uniqueColumns);
        schema.setAttributes(attributes); // 设置属性列表
        return schema;
    }

    private static boolean isPrimaryKeyColumn(SQLColumnDefinition column) {
        return column.getConstraints().stream().anyMatch(
                c -> c instanceof SQLColumnPrimaryKey || "PRIMARY KEY".equalsIgnoreCase(c.toString())
        );
    }

    private static boolean isUniqueColumn(SQLColumnDefinition column) {
        return column.getConstraints().stream().anyMatch(
                c -> c instanceof SQLColumnUniqueKey || "UNIQUE".equalsIgnoreCase(c.toString())
        );
    }

    private static String extractColumnName(SQLExpr expr) {
        if (expr instanceof SQLIdentifierExpr) {
            return ((SQLIdentifierExpr) expr).getName();
        } else if (expr instanceof SQLPropertyExpr) {
            return ((SQLPropertyExpr) expr).getName();
        }
        return expr.toString();  // 其他情况返回字符串表示
    }

    private static void printSchemas(List<TableSchema> schemas) {
        for (TableSchema schema : schemas) {
            System.out.println("\n表名: " + schema.getName());
            System.out.println("主键列: " + schema.getPrimaryKeyColumns());

           // System.out.println("外键列:");
            for (ForeignKeyInfo fk : schema.getForeignKeyColumns()) {
                System.out.printf("  %s → %s(%s)\n",
                        fk.getColumnName(),
                        fk.getReferencedTableName(),
                        fk.getReferencedColumnName()
                );
            }

            //System.out.println("唯一列: " + schema.getUniqueColumns());
            System.out.println("属性列: " + schema.getAttributes()); // 输出属性列
        }
        System.out.println("\n解析完成");
    }
}
